import torch
import torch.optim as optim
import torch.nn.functional as F
from torchvision import datasets, transforms
import gradio as gr
from model import MLP

# 加载MNIST数据集
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])
train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('./data', train=False, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=False)

# 训练模型
model = MLP()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)

def train(epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        
        if batch_idx % 100 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))

for epoch in range(1, 21):  # 增加训练迭代次数
    train(epoch)

# 保存训练好的模型
torch.save(model.state_dict(), 'mnist_model.pt')

# 创建预测函数
def predict(image):
    image = image.reshape((28, 28))
    image = (255 - image).astype('uint8')
    image_tensor = torch.tensor(image).unsqueeze(0).unsqueeze(0).float() / 255.0
    output = model(image_tensor)
    _, predicted = torch.max(output.data, 1)
    return str(predicted.item())

# 加载保存的模型权重
model.load_state_dict(torch.load('mnist_model.pt'))

# 创建Gradio界面
iface = gr.Interface(fn=predict, inputs="sketchpad", outputs="label")

# 启动Gradio界面
iface.launch()